import os
import image_detect
from lxml import etree
import tensorflow as tf
import image_classification
from object_detection.utils import dataset_util
'''使用说明：
image_path：测试图片路径
xml_path：测试图片的标准
outpath：对预测错误的图片结果进行输出
注意：测试图片需要是以一图一车的形式'''
image_path = 'your input image path'
xml_path = 'your input xmls path'
outpath = 'your output path'
imagelist = os.listdir(image_path)
image_detect.init()
image_classification.init()
pass_count = 0
fail_count = 0
'''计算预测结果与标注的IOU'''
def cal_iou(box, ret):
    areas1 = (box[3] - box[1]) * (box[2] - box[0])
    areas2 = (ret[3] - ret[1]) * (ret[2] - ret[0])
    xx1 = max(box[0], ret[0])
    yy1 = max(box[1], ret[1])
    xx2 = min(box[2], ret[2])
    yy2 = min(box[3], ret[3])

    w = max(0.0, xx2 - xx1)
    h = max(0.0, yy2 - yy1)
    inter = w * h
    ret = inter / (areas1 + areas2 - inter)
    print(ret)
    return ret
for imagefile in imagelist:
    ret = image_detect.detecter(os.path.join(image_path, imagefile), 'ssd', outpath)
    print(ret)
    #image_detect.save_result(ret)
    fname, fename = os.path.splitext(imagefile)
    xmlfile = fname + '.xml'
    xml_file = os.path.join(xml_path, xmlfile)
    with tf.gfile.GFile(xml_file, 'r') as fid:
        xml_str = fid.read()
    xml = etree.fromstring(xml_str)

    data = dataset_util.recursive_parse_xml_to_dict(xml)['annotation']
    width = int(data['size']['width'])
    height = int(data['size']['height'])
    class_name = ''
    box = []
    for obj in data['object']:
        box.append(float(obj['bndbox']['ymin']) / height)
        box.append(float(obj['bndbox']['xmin']) / width)
        box.append(float(obj['bndbox']['ymax']) / height)
        box.append(float(obj['bndbox']['xmax']) / width)
        class_name = obj['name']
        break
    if len(ret[os.path.join(image_path, imagefile)]) == 0:
        fail_count = fail_count + 1
        print('detect fail %d' % fail_count)
        print(os.path.join(image_path, imagefile))
        image_detect.save_result(ret)
        continue
    print('label:%d,ret:%d' % (int(class_name.split('Label')[1]), int(ret[os.path.join(image_path, imagefile)][0])))
    if int(ret[os.path.join(image_path, imagefile)][0]) == int(class_name.split('Label')[1]):
        if cal_iou(box, ret[os.path.join(image_path, imagefile)][2]) > 0.85:
            pass_count = pass_count + 1
            print('pass %d' % pass_count)
        else:
            fail_count = fail_count + 1
            print('local fail %d' % fail_count)
            print(os.path.join(image_path, imagefile))
            image_detect.save_result(ret)
    else:
        fail_count = fail_count + 1
        print('class fail %d' % fail_count)
        print(os.path.join(image_path, imagefile))
        image_detect.save_result(ret)
print('img num:%d,pass num:%d,fail num:%d,accuracy:%f' % (len(imagelist), pass_count, fail_count, pass_count/len(imagelist)))
